
# 该程序将mujoco的xml文件转换为栅格地图，仅能处理mujoco中的简单障碍物格式：box, cylinder

import xml.etree.ElementTree as ET
from pyquaternion import Quaternion
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.patches as patches
from matplotlib.transforms import Affine2D
import math
import os
import json
import cv2
from io import BytesIO
import pickle
import io

class MJCF2GRID:
    def __init__(
        self, 
        scene_dir, 
        collision_group="1",
        safe_distance: int = 0.5,
        resolution=0.01
    ):
        scene_dir_name = os.path.basename(scene_dir)
        self.mjcf_path = os.path.join(scene_dir, f"{scene_dir_name}.xml")
        self.output_image = os.path.join(scene_dir, f"{scene_dir_name}.png")
        params_path = os.path.join(scene_dir, f"{scene_dir_name}.json")
        
        self.safe_distance = safe_distance
        self.params= {}
        if os.path.exists(params_path):
            with open(params_path, "r") as f:
                self.params = json.load(f)
        self.collision_group = collision_group
        self.resolution = resolution
        
    def run(self):
        # 读取xml文件
        tree = ET.parse(self.mjcf_path)
        # 获取根节点
        root = tree.getroot()
        # 获取地图尺寸
        worldbody = root.find('worldbody')
        floor_size = self.get_floor_size(worldbody)
        collision_info = self.get_collision_info(worldbody)
        self.plt_draw(floor_size, collision_info, self.output_image)

    # 找到floor字段，获取尺寸
    def get_floor_size(self, worldbody):
        for geom in worldbody.findall('geom'):
            if geom.get('name') == 'floor':
                size = geom.get('size')
                size = [float(x) for x in size.split()]
                return size

    def get_collision_info(self, worldbody):
        def quat_to_xy_angle(quat):
            q = Quaternion(quat)
            euler = q.yaw_pitch_roll
            return euler[0]  # Z-axis rotation

        def check_if_collision(geom):
            ignore = ["floor", "hfield"]
            if geom.get("class") == "site":
                return False
            
            if geom.get('name') not in ignore:
                if geom.get("group") is None:
                    return True
                elif geom.get("group") == self.collision_group:
                    return True
            return False
                
        # 存储geom信息的字典
        collision_info = []       
        for geom in worldbody.findall('geom'):
            if check_if_collision(geom):
                info = {
                    'type': geom.get('type'),
                    'size': [float(x) for x in geom.get('size').split()],
                    # 'size': geom.get('size'),
                    'pos': [float(x) for x in geom.get('pos').split()],
                    'yaw': None
                }
                quat_str = geom.get('quat')
                if quat_str:

                    quat = [float(x) for x in quat_str.split()]
                    info['yaw'] = quat_to_xy_angle(quat)
                collision_info.append(info)
        return collision_info

    def plt_draw(self, floor_size, collision_info, output_image):
        dpi_figure = 1 / self.resolution
        # 创建一个分辨率为 1000x1000 的画布
        fig, ax = plt.subplots(figsize=(floor_size[0] * 2, floor_size[1] * 2), dpi=dpi_figure)  # figsize (width, height) in inches, dpi=100
                                            # 10 inches * 100 dpi = 1000 pixels in each dimension                               
        for collision in collision_info:
            if collision['type'] == 'box':
                # 创建一个矩形
                rect = patches.Rectangle(
                    (
                        (collision['pos'][0] + floor_size[0] - collision['size'][0]) * dpi_figure, 
                        (collision['pos'][1] + floor_size[1] - collision['size'][1]) * dpi_figure
                    ),  # 矩形左下角的坐标
                    (collision['size'][0]) * 2 * dpi_figure,         # 矩形的宽度
                    (collision['size'][1]) * 2 * dpi_figure,         # 矩形的高度
                    linewidth=0, 
                    facecolor='black'  # 填充颜色
                )
                if collision['yaw'] != None:
                    # 创建一个旋转矩阵
                    rotation_center = ((collision['pos'][0] + floor_size[0]) * dpi_figure, 
                                    (collision['pos'][1] + floor_size[1]) * dpi_figure)  # 矩形的中心
                    angle_deg = (collision['yaw'] / math.pi) * 180
                    rotation = Affine2D().rotate_deg_around(rotation_center[0], rotation_center[1], angle_deg)
                    rect.set_transform(rotation + ax.transData)
                ax.add_patch(rect)
            else:
                # 创建一个圆形
                circle = patches.Circle(
                    ((collision['pos'][0] + floor_size[0]) * dpi_figure, 
                    (collision['pos'][1] + floor_size[1]) * dpi_figure),   # 圆心位置
                    collision['size'][0] * dpi_figure,          # 半径
                    linewidth=0,
                    facecolor='black'  # 填充颜色，黑色
                )
                ax.add_patch(circle)
    

        # 设置坐标轴范围
        ax.set_xlim(0, floor_size[0] * 2 * dpi_figure)
        ax.set_ylim(0, floor_size[1] * 2 * dpi_figure)
        # 显示画布
        plt.axis('off')  # 隐藏坐标轴
        safe_mask = self.gen_safe_area(fig, floor_size, dpi_figure)
        
        buf = BytesIO()
        fig.savefig(buf, format='png', dpi=dpi_figure, bbox_inches='tight', pad_inches=0)
        img_array = np.frombuffer(buf.getvalue(), dtype=np.uint8)
        img_cv = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
        resized_img = cv2.resize(
            img_cv, 
            (
                int(floor_size[0] * 2 * dpi_figure), 
                int(floor_size[1] * 2 * dpi_figure)
            )
        )
        gray_img = cv2.cvtColor(resized_img, cv2.COLOR_BGR2GRAY)  # 转为灰度
        _, binary_img = cv2.threshold(gray_img, 127, 255, cv2.THRESH_BINARY)  # 二值化
        cv2.imwrite(output_image, binary_img)
        
        safe_area_output = output_image.replace('.png', '_safe.png')
        # 创建彩色图片以标记安全区域
        color_image = cv2.cvtColor(binary_img, cv2.COLOR_GRAY2BGR)
        # 将安全区域标记为绿色
        color_image[safe_mask > 0] = [0, 255, 0]  # 绿色
        cv2.imwrite(safe_area_output, color_image)
        
        plt.close(fig)


    def gen_safe_area(self, fig, floor_size, dpi_figure):
        # 使用pickle复制
        buf = io.BytesIO()
        pickle.dump(fig, buf)
        buf.seek(0)
        fig_copy = pickle.load(buf)
        # 获取新图形的axes
        ax_copy = fig_copy.axes[0]

        if self.params.get("mode", "random") == "stairwell":
            rect = patches.Rectangle(
                (
                    0,
                    (floor_size[1] + self.params["wall"]["center_aisle_with"] / 2) * dpi_figure
                ),
                floor_size[0] * 2 * dpi_figure, 
                (floor_size[1] - self.params["wall"]["center_aisle_with"] / 2) * dpi_figure,
                linewidth=0, 
                facecolor='black'
            )
            ax_copy.add_patch(rect)
            rect = patches.Rectangle(
                (0, 0),
                (floor_size[0] - self.params["wall"]["side_aisle_with"] / 2) * dpi_figure, 
                (floor_size[1] - self.params["wall"]["center_aisle_with"] / 2) * dpi_figure,
                linewidth=0, 
                facecolor='black'
            )
            ax_copy.add_patch(rect)
            rect = patches.Rectangle(
                (
                    (floor_size[0] + self.params["wall"]["side_aisle_with"] / 2) * dpi_figure, 
                    0
                ),
                (floor_size[0] - self.params["wall"]["side_aisle_with"] / 2) * dpi_figure,
                (floor_size[1] - self.params["wall"]["center_aisle_with"] / 2) * dpi_figure,
                linewidth=0, 
                facecolor='black'
            )
            ax_copy.add_patch(rect)
            
        buf = BytesIO()
        fig_copy.savefig(buf, format='png', dpi=dpi_figure, bbox_inches='tight', pad_inches=0)
        img_array = np.frombuffer(buf.getvalue(), dtype=np.uint8)
        img_cv = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
        resized_img = cv2.resize(
            img_cv, 
            (
                int(floor_size[0] * 2 * dpi_figure), 
                int(floor_size[1] * 2 * dpi_figure)
            )
        )
        binary_image = resized_img.astype(np.uint8)
        binary_image = cv2.cvtColor(binary_image, cv2.COLOR_BGR2GRAY)
        _, binary_image = cv2.threshold(binary_image, 127, 255, cv2.THRESH_BINARY)
        # 距离变换，计算每个白色像素到最近黑色像素的距离
        distance = cv2.distanceTransform(binary_image, cv2.DIST_L2, cv2.DIST_MASK_5)
        # 创建掩码：距离大于等于safe_distance_pixel的区域为安全区域
        safe_distance_pixel = self.safe_distance / self.resolution
        safe_mask = (distance >= safe_distance_pixel).astype(np.uint8) * 255
        plt.close(fig_copy)
        return safe_mask
        

        
if __name__ == '__main__':
    
    scenes_dir = "dataset_tools/sources/mjcf_scenes"
    
    scenes_dir_list = [
        dir 
        for dir in os.listdir(scenes_dir) 
        if os.path.isdir(os.path.join(scenes_dir, dir)) and "scene_" in dir
    ]
    for scene_dir in scenes_dir_list:
        mjcf2grid = MJCF2GRID(os.path.join(scenes_dir, scene_dir), resolution=0.02)
        mjcf2grid.run()
    
    
